﻿using System;
using System.Collections.Generic;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;

namespace Mars.GpuInstance
{
    public class InstanceRenderBatch : IRecyclable,IDisposable
    {
        protected static Matrix4x4[] matrices = new Matrix4x4[InstancerConstants.COUNT_PER_DRAW];
        protected static Vector4[] campColorDatas = new Vector4[InstancerConstants.COUNT_PER_DRAW];
        
        public MaterialPropertyBlock mpb;
        protected NativeArray<int> outputCount;
        public NativeArray<float4x4> output;
        public NativeArray<float4> campData;
        public Matrix4x4 offsetMatrix;
        
        public Mesh meshItem;
        public List<Material> materialItems = new();

        private Dictionary<string, bool> keywordDic = new();

        //渲染范围，只有距离相机在这个范围内的才会渲染
        private float2 drawRange;
        public int lodLevel = 0;

        private bool hasOffsetMatrix;

        public InstanceRenderBatch()
        {
            outputCount = new NativeArray<int>(1, Allocator.Persistent);
        }
        
        public void UpdateDrawRange(float2 r)
        {
            drawRange = r;
        }

        public void UpdateLodLevel(int lod)
        {
            lodLevel = lod;
        }

        public void UpdateOffsetMatrix(Matrix4x4 matrix)
        {
            offsetMatrix = matrix;
            hasOffsetMatrix = !offsetMatrix.Equals(Matrix4x4.identity);
        }

        public void UpdateMpb(MaterialPropertyBlock m)
        {
            mpb = m;
        }

        public void UpdateRenderAsset(Mesh mesh, List<Material> _materials)
        {
            meshItem = mesh;
            materialItems.Clear();
            materialItems.AddRange(_materials);
        }

        public void EnableKeyWord(string keyWord)
        {
            if (keywordDic.ContainsKey(keyWord))
            {
                keywordDic[keyWord] = true;
            }
            else
            {
                keywordDic.Add(keyWord, true);
            }
            foreach (var material in materialItems)
            {
                if (material)
                {
                    material.EnableKeyword(keyWord);
                }
            }
        }

        public void DisableKeyWord(string keyWord)
        {
            if (keywordDic.ContainsKey(keyWord))
            {
                keywordDic[keyWord] = false;
            }
            else
            {
                keywordDic.Add(keyWord, false);
            }
            
            foreach (var material in materialItems)
            {
                if (material)
                {
                    material.DisableKeyword(keyWord);
                }
            }
        }

        public int GetInstanceCount()
        {
            if (outputCount.IsCreated)
            {
                return outputCount[0];
            }

            return 0;
        }

        public void ClearDrawCount()
        {
            outputCount[0] = 0;
        }

        public JobHandle CopyData(InstanceRuntimeData runtimeData, bool isShadow, JobHandle dependsOn)
        {
            var len = runtimeData.input.Length;

            if (len <= 0)
            {
                outputCount[0] = 0;
                return default;
            }
            
            if (!output.IsCreated)
            {
                output = new NativeArray<float4x4>(len, Allocator.Persistent);
            }
            else if (output.IsCreated && output.Length < len)
            {
                output.Dispose();
                output = new NativeArray<float4x4>(len, Allocator.Persistent);
            }

            if (!campData.IsCreated)
            {
                campData = new NativeArray<float4>(len, Allocator.Persistent);
            }
            else if (campData.IsCreated && campData.Length < len)
            {
                campData.Dispose();
                campData = new NativeArray<float4>(len, Allocator.Persistent);
            }

            outputCount[0] = 0;

            var job = new CopyDataJob
            {
                input = runtimeData.input,
                viewIndex = runtimeData.viewIndex,
                lodInfo = runtimeData.lodInfo,
                inputCampData = runtimeData.campColorData,
                lod = lodLevel,
                isShadow = isShadow,
                output = output,
                outputCount = outputCount,
                campData = campData,
                hasOffsetMatrix = hasOffsetMatrix,
                offsetMatrix = offsetMatrix
            };

            return job.Schedule(dependsOn);
        }


        #region 绘制相关

        public virtual void Draw(ShadowCastingMode shadowCastingMode, bool hasCampData)
        {
            Draw(GetInstanceCount(), shadowCastingMode, hasCampData, meshItem, materialItems, output, campData, mpb);
        }

        public static void Draw(int instanceCount, ShadowCastingMode shadowCastingMode, bool hasCampData,
            Mesh mesh, 
            List<Material> drawMaterials,
            NativeArray<float4x4> allMatrices,
            NativeArray<float4> allCampData,
            MaterialPropertyBlock mpb)
        {
            if (instanceCount <= 0 || mesh == null)
            {
                return;
            }

            if (!allMatrices.IsCreated || !allCampData.IsCreated)
            {
                return;
            }
            

            for (int j = 0; j < drawMaterials.Count; j++)
            {
                var drawMat = drawMaterials[j];

                if (drawMat == null) continue;

                var submeshIndex = Math.Min(j, mesh.subMeshCount - 1);

                if (instanceCount > InstancerConstants.COUNT_PER_DRAW) //要分成多个批次画
                {
                    var batchCount = (int) math.ceil((float) instanceCount / (float) InstancerConstants.COUNT_PER_DRAW);
                    for (int k = 0; k < batchCount; k++)
                    {
                        var srcIndex = k * InstancerConstants.COUNT_PER_DRAW;
                        var copyCount = InstancerConstants.COUNT_PER_DRAW;

                        if (k == (batchCount - 1))
                        {
                            copyCount = instanceCount - k * InstancerConstants.COUNT_PER_DRAW;
                        }

                        if (hasCampData)
                        {
                            MathUtil.Copy<float4, Vector4>(allCampData, srcIndex, campColorDatas, 0, copyCount);
                        }
                        MathUtil.Copy<float4x4, Matrix4x4>(allMatrices, srcIndex, matrices, 0, copyCount);
                        if (hasCampData)
                        {
                            mpb.SetVectorArray(InstancerConstants.GPU_CampColorData, campColorDatas);
                        }

                        Graphics.DrawMeshInstanced(mesh, submeshIndex, drawMat, matrices, copyCount, mpb, shadowCastingMode);
                    }
                }
                else
                {
                    MathUtil.Copy<float4x4, Matrix4x4>(allMatrices, 0, matrices, 0, instanceCount);
                    if (hasCampData)
                    {
                        MathUtil.Copy<float4, Vector4>(allCampData, 0, campColorDatas, 0, instanceCount);
                    }
                    if (hasCampData)
                    {
                        mpb.SetVectorArray(InstancerConstants.GPU_CampColorData, campColorDatas);
                    }

                    Graphics.DrawMeshInstanced(mesh, submeshIndex, drawMat, matrices, instanceCount, mpb, shadowCastingMode);
                }
            }
        }
        #endregion


        #region 释放回收相关

        public virtual void Dispose()
        {
            if (outputCount.IsCreated)
            {
                outputCount.Dispose();
            }
            
            if (output.IsCreated)
            {
                output.Dispose();
            }

            if (campData.IsCreated)
            {
                campData.Dispose();
            }

            keywordDic.Clear();
            mpb = null;
            meshItem = null;
            materialItems.Clear();
        }

        public virtual void OnRecycle()
        {
            if (output.IsCreated)
            {
                output.Dispose();
            }

            if (campData.IsCreated)
            {
                campData.Dispose();
            }

            keywordDic.Clear();
            
            mpb = null;
            meshItem = null;
            materialItems.Clear();
        }
        #endregion
    }
}